"""
深度残差网络：CIFAR10分类
"""

import os
import torch
import torch.nn as nn
import torchvision.datasets
import torch.utils.data
import torchvision.transforms as transforms
from tqdm import tqdm
import numpy as np
import matplotlib.pyplot as plt


def conv3x3(in_channels, out_channels, stride=1):
    """卷积"""
    return nn.Conv2d(in_channels, out_channels, kernel_size=(3, 3), stride=(stride, stride), padding=1, bias=False)


class ResidualBlock(nn.Module):
    """残差模块"""

    def __init__(self, in_channels, out_channels, stride=1, down_sample=None):
        super().__init__()
        self.conv1 = conv3x3(in_channels, out_channels, stride)
        self.bn1 = nn.BatchNorm2d(out_channels)
        self.relu = nn.ReLU(inplace=True)
        self.conv2 = conv3x3(out_channels, out_channels)
        self.bn2 = nn.BatchNorm2d(out_channels)
        self.down_sample = down_sample

    def forward(self, x):
        residual = x
        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)
        out = self.conv2(out)
        out = self.bn2(out)
        if self.down_sample:
            residual = self.down_sample(x)
        out += residual
        self.relu(out)
        return out


class ResNet(nn.Module):
    """残差网络"""

    def __init__(self, block, layers, num_classes=10):
        super().__init__()
        self.in_channels = 16
        self.conv = conv3x3(3, self.in_channels)
        self.bn = nn.BatchNorm2d(self.in_channels)
        self.relu = nn.ReLU(inplace=True)
        self.layer1 = self.make_layer(block, 16, layers[0])
        self.layer2 = self.make_layer(block, 32, layers[1], 2)
        self.layer3 = self.make_layer(block, 64, layers[2], 2)
        self.avg_pool = nn.AvgPool2d(8)
        self.fc = nn.Linear(64, num_classes)

    def make_layer(self, block, out_channels, blocks, stride=1):
        down_sample = None
        if stride != 1 or self.in_channels != out_channels:
            down_sample = nn.Sequential(conv3x3(self.in_channels, out_channels, stride=stride),
                                        nn.BatchNorm2d(out_channels))
        layers = [block(self.in_channels, out_channels, stride, down_sample)]
        self.in_channels = out_channels
        for _ in range(1, blocks):
            layers.append(block(out_channels, out_channels))
        return nn.Sequential(*layers)

    def forward(self, x):
        out = self.conv(x)
        out = self.bn(out)
        out = self.relu(out)
        out = self.layer1(out)
        out = self.layer2(out)
        out = self.layer3(out)
        out = self.avg_pool(out)
        out = out.view(out.size(0), -1)
        out = self.fc(out)
        return out


def update_lr(optimizer, lr):
    """更新学习率"""
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr


class Config(object):
    """单例模式，配置参数"""
    _instance = None
    num_classes = 10
    batch_size = 4
    num_epochs = 100
    nw = os.cpu_count()
    learning_rate = 0.001
    model_path = 'model/checkpoint/resnet-1.pth'
    device_cuda = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    device_cpu = torch.device('cpu')

    def __new__(cls, *args, **kw):
        if cls._instance is None:
            cls._instance = object.__new__(cls, *args, **kw)
        return cls._instance


class Cifar10:
    _instance = None
    config = Config()
    batch_size = config.batch_size
    nw = config.nw

    transform = transforms.Compose([transforms.Pad(4),
                                    transforms.RandomHorizontalFlip(),
                                    transforms.RandomCrop(32),
                                    transforms.ToTensor()])

    trainset = torchvision.datasets.CIFAR10(root='./dataset/cifar10', train=True, download=True, transform=transform)
    trainloader = torch.utils.data.DataLoader(dataset=trainset, batch_size=config.batch_size, shuffle=True,
                                              num_workers=config.nw)

    testset = torchvision.datasets.MNIST(root='./dataset/cifar10', train=False, download=True, transform=transform)
    testloader = torch.utils.data.DataLoader(dataset=testset, batch_size=batch_size, shuffle=False, num_workers=nw)

    def __new__(cls, *args, **kw):
        if cls._instance is None:
            cls._instance = object.__new__(cls, *args, **kw)
        return cls._instance


def train():
    """训练模型"""
    config = Config()
    device = config.device_cuda

    model_path = config.model_path
    model = ResNet(ResidualBlock, [2, 2, 2]).to(device)
    if os.path.isfile(model_path):
        print('load model from: ', model_path)
        model.load_state_dict(torch.load(model_path))
        model.train()

    optimizer = torch.optim.Adam(model.parameters(), lr=config.learning_rate)
    criterion = nn.CrossEntropyLoss()
    trainloader = Cifar10().trainloader
    curr_lr = config.learning_rate
    for epoch in range(config.num_epochs):
        print(f'Epoch:{epoch + 1}/{config.num_epochs}')
        pbar = tqdm(trainloader)
        running_loss = []
        for images, labels in pbar:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            running_loss.append(loss.item())
            pbar.set_description_str(f'loss:{round(np.nanmean(running_loss), 3)}')

        if (epoch + 1) % 20 == 0:
            curr_lr /= 3
            update_lr(optimizer, curr_lr)

    # save model
    print(f'save model to: {model_path}')
    if not os.path.isdir(os.path.dirname(model_path)):
        os.makedirs(os.path.dirname(model_path))
    torch.save(model.state_dict(), model_path)


def test():
    """测试模型"""
    config = Config()
    model_path = config.model_path
    device = config.device_cuda
    model = ResNet(ResidualBlock, [2, 2, 2]).to(device)
    if os.path.isfile(model_path):
        print('load model from: ', model_path)
        model.load_state_dict(torch.load(model_path))

    model.eval()
    testloader = Cifar10().testloader
    with torch.no_grad():
        correct = 0
        total = 0
        for images, labels in testloader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            _, preds = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (preds == labels).sum().item()
        print(f'Test accuracy on {total} examples:{round(correct / total * 100, 2)}%')


def view():
    """可视化显示样本"""
    config = Config()
    model_path = config.model_path
    device = config.device_cpu
    model = ResNet(ResidualBlock, [2, 2, 2]).to(device)
    if os.path.isfile(model_path):
        print('load model from: ', model_path)
        model.load_state_dict(torch.load(model_path))

    model.eval()
    model.to(device)
    testloader = Cifar10().testloader
    images, labels = next(iter(testloader))
    with torch.no_grad():
        outputs = model(images)
        _, preds = torch.max(outputs, 1)
    plt.figure()
    for i, image in enumerate(images.numpy()[:16, :, :]):
        plt.subplot(4, 4, i + 1)
        plt.imshow(image.squeeze())
        plt.title(f'label:{labels[i].item()},pred:{preds[i].item()}')
        plt.axis('off')
    plt.show()


if __name__ == '__main__':
    train()
    test()
    view()
